package com.fjl.list.leetcode373;
import java.util.ArrayList;
import java.util.List;
import java.util.PriorityQueue;

/**
 * @Description
 * @Author: 01195
 * @Date: 2025/2/19 15:01
 */
//给定两个以 非递减顺序排列 的整数数组 nums1 和 nums2 , 以及一个整数 k 。
//定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。请找到和最小的 k 个数对
//输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
//输出: [1,2],[1,4],[1,6]
//解释: 返回序列中的前 3 对数：
//[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
public class FindKSmallestPairs {
    //这道题其实是前文 单链表的六大解题套路 中讲过的 23. 合并K个升序链表 的变体。
    //怎么把这道题变成合并多个有序链表呢？就比如说题目输入的用例：
    //nums1 = [1,7,11], nums2 = [2,4,6]
    //组合出的所有数对儿这就可以抽象成三个有序链表：
    //[1, 2] -> [1, 4] -> [1, 6]
    //[7, 2] -> [7, 4] -> [7, 6]
    //[11, 2] -> [11, 4] -> [11, 6]
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // 存储三元组 (num1[i], nums2[i], i)
        // i 记录 nums2 元素的索引位置，用于生成下一个节点
        PriorityQueue<int[]> queue = new PriorityQueue<>(
                (a, b) -> {
                    // 按照数对的元素和升序排序
                    return (a[0] + a[1]) - (b[0] + b[1]);
                });
        // 按照 23 题的逻辑初始化优先级队列
        for (int i = 0; i < nums1.length; i++) {
            queue.add(new int[]{nums1[i], nums2[0], 0});
        }

        List<List<Integer>> result = new ArrayList<>();
        while (!queue.isEmpty() && k > 0) {
            int[] poll = queue.poll();
            k--;
            int pos = poll[2];
            // 链表中的下一个节点加入优先级队列
            if (pos + 1 < nums2.length) {
                queue.add(new int[]{poll[0], nums2[pos + 1], pos + 1});
            }
            List<Integer> list = new ArrayList<>();
            list.add(poll[0]);
            list.add(poll[1]);
            result.add(list);
        }
        return result;
    }
}
